#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot shallow water Rossby-Haurwitz wave test results')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-o', '--output', help='Output figure')
args = parser.parse_args()

f = xr.open_dataset('bw.1440x720x30.h0.nc')

proj = ccrs.PlateCarree(central_longitude=180)
fig = plt.figure(figsize=(8, 4))

days = ((f.time[args.time_step] - f.time[0]).values).astype('timedelta64[D]')

fig.suptitle(days, fontsize=12)

if days == 7:
	levels = np.linspace(98400, 100800, 7)
	colors = None
elif days == 9:
	levels = np.linspace(96400, 101600, 14)
	colors = ListedColormap([
		[0.07, 0.16, 0.96],
		[0.09, 0.49, 0.96],
		[0.15, 0.98, 0.99],
		[0.13, 0.84, 0.55],
		[0.11, 0.69, 0.20],
		[0.40, 0.79, 0.21],
		[1.00, 0.99, 0.27],
		[1.00, 0.78, 0.23],
		[0.99, 0.60, 0.19],
		[0.99, 0.35, 0.15],
		[0.99, 0.00, 0.11],
		[0.61, 0.00, 0.24],
		[0.31, 0.05, 0.45]
	])
else:
	levels = None
	colors = None

ax = fig.add_subplot(1, 1, 1, projection=proj)
plot_contour_map(ax, f.phs[args.time_step,:,:], cmap=colors, left_string='PS', levels=levels)

if args.output:
	plt.savefig(args.output)
else:
	plt.show()
